
#include "project.h"

static char *keyfile = NULL;
enum action {
	QUIT,
	REG,
	TALLY_KEY,
	VOTE,
	TALLY
};
static int action = QUIT;
static BN_CTX *bn_ctx = NULL;
// CTF's public RSA key
static RSA *CTFPub_RSA = NULL;
// Voter's RSA key (pub/priv)
static RSA *My_RSA = NULL;
// Voter's public RSA key
static RSA *MyPub_RSA = NULL;
// Voter's private RSA key
static RSA *MyPriv_RSA = NULL;
static int Dv_len = 0;
static unsigned char *Dv = NULL;
static int Ev_len = 0;
static unsigned char *Ev = NULL;
static unsigned char hEv[SHA_DIGEST_LENGTH];
static unsigned char *DchEv= NULL;
static int size_DchEv = 0;
// Voter's ballot
static unsigned long B[2];
static bool K_valid = false;
// Voter's secret key used to encrypt the ballot
static unsigned char K[16];
// Voter's username
static int uid = -1;
// Voter's password
static char *passwd = NULL;

// read keyfile
void
read_keys(void)
{
	FILE *s = fopen(keyfile, "rb");
	do {
		// read voter's public key (tallying key), Ev
		// we read the length first and then that many bytes
		if (fread(&Ev_len, sizeof(int), 1, s) != 1)
			break;
		Ev = (unsigned char *)malloc(Ev_len);
		if (fread(Ev, Ev_len, 1, s) != 1)
			break;
		// read voter's private key Dv
		if (fread(&Dv_len, sizeof(int), 1, s) != 1)
			break;
		Dv = (unsigned char *)malloc(Dv_len);
		if (fread(Dv, Dv_len, 1, s) != 1)
			break;
		// read CTF's signature Dc(h(Ev))
		if (fread(&size_DchEv, sizeof(int), 1, s) != 1)
			break;
		DchEv = (unsigned char *)malloc(size_DchEv);
		if (fread(DchEv, size_DchEv, 1, s) != 1)
			break;
		// read secret key K. this field is only present in the
		// keyfile after voter has casted his ballot at least once.
		// it is always the last secret key he used.
		unsigned char temp[16];
		if (fread(&temp, 16, 1, s) == 1){
			memmove(K, temp, 16);
			K_valid = true;
		}
		else {
			memset(K, 0, 16);
			K_valid = false;
		}
		fclose(s);
		// create RSA structs for use with openssl crypto library
		// from data read in from the file.
		unsigned char *Key = Ev;
		MyPub_RSA = d2i_RSAPublicKey(NULL, (const unsigned char **)&Key, Ev_len);
		Key = Dv;
		MyPriv_RSA = d2i_RSAPrivateKey(NULL, (const unsigned char **)&Key, Dv_len);
		assert(MyPub_RSA && MyPriv_RSA);
		return;
	} while (0);
	fclose(s);
	printf("ERROR reading keyfile.\n");
	abort();
}

// write keyfile
void
write_keys(void)
{
	FILE *s = fopen(keyfile, "wb");
	do {
		// write Ev
		if (fwrite(&Ev_len, sizeof(int), 1, s) != 1)
			break;
		if (fwrite(Ev, Ev_len, 1, s) != 1)
			break;
		// write Dv
		if (fwrite(&Dv_len, sizeof(int), 1, s) != 1)
			break;
		if (fwrite(Dv, Dv_len, 1, s) != 1)
			break;
		// write CTF's signature
		if (fwrite(&size_DchEv, sizeof(int), 1, s) != 1)
			break;
		if (fwrite(DchEv, size_DchEv, 1, s) != 1)
			break;
		// if we have voted, write secret key K
		if (K_valid && fwrite(K, 16, 1, s) != 1)
			break;
		fclose(s);
		return;
	} while (0);
	fclose(s);
	printf("ERROR writing keyfile.\n");
	abort();
}



// contact CTF and obtain its public key
void
get_CTF_Key(void)
{
	MSG m;
	int data;
	m.size = 4;
	m.data = &data;
	MSG_TYPE(&m) = MSG_GET_CTF_PUBLIC_KEY;
	if (send_recv(&m, &m) == 0){
		CTFPub_RSA = get_ctf_public_key(&m);
		free(m.data);
		if (CTFPub_RSA)
			return;
	}
	printf("ERROR: get_CTF_Key.\n");
	assert(0);
}

// register to vote
void
register_to_vote(void)
{
	// check whether we have valid username, password
	if (passwd == NULL || uid < 0){
		printf("ERROR! uid or passwd or both not set.\n");
		abort();
	}

	// compute an RSA key pair with modulus greater than that of CTF
	// CTF's modulus has 1024 bits. So pick a modulus having 1280 bits
	// we check to make sure that modulus is indeed greater than that of CTF's
	do {
		if (My_RSA)
			RSA_free(My_RSA);
		My_RSA = RSA_generate_key(1280, 17, NULL, NULL);
	} while (BN_cmp(My_RSA->n, CTFPub_RSA->n) <= 0);

	// convert voters public key (Ev) to an array of bytes to be sent to CTF
	Ev_len = i2d_RSAPublicKey(My_RSA, 0);
	Ev = (unsigned char *)malloc(Ev_len);
	unsigned char *Key = Ev;
	assert(i2d_RSAPublicKey(My_RSA, &Key) == Ev_len);

	// compute h(Ev), use SHA1 hash
	assert(SHA1(Ev, Ev_len, hEv));
	// X is h(Ev)
	BIGNUM *X = BN_bin2bn(hEv, sizeof(hEv), NULL);
	BIGNUM *R = BN_new();
	// choose an odd random number R with 160 bits, with most significant bit 1
	assert(BN_rand(R, 160, 1, 1) == 1);
	// R^(-1) mod N where N is the CTF's modulus
	BIGNUM *RI = BN_mod_inverse(NULL, R, CTFPub_RSA->n, bn_ctx);
	BIGNUM *R1 = BN_new();
	// Ec(R)
	assert(BN_mod_exp(R1, R, CTFPub_RSA->e, CTFPub_RSA->n, bn_ctx) == 1);
	// R1 is Ec(R) * h(Ev)
	assert(BN_mod_mul(R1, R1, X, CTFPub_RSA->n, bn_ctx) == 1);

	// convert R1 to bytes, buf1 is Ec(R) * h(Ev)
	int len1 = BN_num_bytes(R1);
	unsigned char buf1[len1];
	assert(BN_bn2bin(R1, buf1) == len1);

	// this part computes Ec(h(buf1), username, password)
	int len2 = strlen(passwd);
	int size = SHA_DIGEST_LENGTH+4+len2+1+2;


	unsigned char buf2[size];
#define	BUFP(type, offset) ((type *)&buf2[offset])
	*BUFP(unsigned char, 0) = 0xff;
	*BUFP(unsigned char, size-1) = 0xff;
	// h(buf1)
	assert(SHA1(buf1, len1, BUFP(unsigned char, 0+1)));
	*BUFP(int, SHA_DIGEST_LENGTH+1) = uid;
	strcpy(BUFP(char, 4+SHA_DIGEST_LENGTH+1), passwd);
	// now buf2 has (h(buf1), username, password)
	BIGNUM *Y = BN_bin2bn(buf2, size, NULL);
	// Y is buf2 as BIGNUM
	BIGNUM *R4 = BN_new();
	// R4 is Ec(h(buf1), username, password)
	assert(BN_mod_exp(R4, Y, CTFPub_RSA->e, CTFPub_RSA->n, bn_ctx) == 1);
	int len3 = BN_num_bytes(R4);

	// create M_1 (see doc.)
	// M_1 is (Ec(R) * h(Ev), Ec(h(buf1), username, password))
	MSG m, r;
	alloc_msg(4+4+len1+4+len3, &m);
	MSG_TYPE(&m) = MSG_BLND_SIGN_VOTING_KEY;
	// copy Ec(R) * h(Ev) into M_1
	*MSG_BODY(&m, 4, int) = BN_bn2bin(R1, MSG_BODY(&m, 8, unsigned char));
	int offset = 8+len1;
	// copy Ec(h(buf1), username, password) into M_1
	*MSG_BODY(&m, offset, int) = BN_bn2bin(R4, MSG_BODY(&m, offset+4, unsigned char));
	// send M_1 to CTF, CTF's response in r
	if (send_recv(&m, &r) == 0){
		assert(MSG_TYPE(&r) == MSG_BLND_SIGN_VOTING_KEY_REPL);
		// R2 has Dc(Ec(R) * h(Ev)) or R * Dc(h(Ev))
		BIGNUM *R2 = BN_bin2bn(MSG_BODY(&r, 8, unsigned char), *MSG_BODY(&r, 4, int), NULL);
		MSG_FREE(&r);
		// now R2 has Dc(h(Ev))
		assert(BN_mod_mul(R2, R2, RI, CTFPub_RSA->n, bn_ctx) == 1);

		BIGNUM *R3 = BN_new();
		// R3 is Ec(Dc(h(Ev))) or h(Ev)
		assert(BN_mod_exp(R3, R2, CTFPub_RSA->e, CTFPub_RSA->n, bn_ctx) == 1);
		// compare R3 with X see above.
		if (BN_cmp(X, R3)){
			printf("ERROR! h(Ev) mismatch.\n");
			abort();
		}

		// save CTF's signature for later. we need to send this back to CTF when we provide
		// CTF our tallying key in func. submit_tally_key()
		size_DchEv = BN_num_bytes(R2);
		DchEv = (unsigned char *)malloc(size_DchEv);
		assert(BN_bn2bin(R2, DchEv) == size_DchEv);

//		BN_print_fp(stdout, R3);
//		printf("\n");
//		BN_print_fp(stdout, X);
//		printf("\n");

		BN_free(R2);
		BN_free(R3);
	}
	else {
		printf("ERROR: 2(c).\n");
		exit(0);
	}
	MSG_FREE(&m);

	BN_free(X);
	BN_free(R);
	BN_free(RI);
	BN_free(R1);
	BN_free(R4);
	BN_free(Y);

	// convert voter's private RSA key to bytes so that we can save it in keyfile
	// we have already converted his public key to bytes when we sent it to CTF (see above)
	Dv_len = i2d_RSAPrivateKey(My_RSA, 0);
	Dv = (unsigned char *)malloc(Dv_len);
	Key = Dv;
	assert(i2d_RSAPrivateKey(My_RSA, &Key) == Dv_len);
}

// submit tallying key to CTF with CTF's signature
void
submit_tally_key(void)
{
	// sign CTF's signature with voter's private key. this require's
	// voter's modulus to be greater than that of CTF's. this is however
	// easily avoided by signing hash of CTF's signature instead of the
	// signature itself. to keep things simple we just do this for now.
	BIGNUM *R2 = BN_bin2bn(DchEv, size_DchEv, NULL);
	BIGNUM *R3 = BN_new();
	// R3 is now Dv(Dc(h(Ev)))
	assert(BN_mod_exp(R3, R2, MyPriv_RSA->d, MyPriv_RSA->n, bn_ctx) == 1);
	int size_DvDchEv = BN_num_bytes(R3);

	// send tallying key to CTF
	// M_3 has the format (Ev, Dv(Dc(h(Ev))))
	MSG m, r;
	alloc_msg(4+4+Ev_len+4+size_DvDchEv, &m);
	MSG_TYPE(&m) = MSG_VOTING_KEY;
	int offset = 4;
	*MSG_BODY(&m, offset, int) = Ev_len;
	offset += sizeof(int);
	// copy Ev into M_3
	memmove(MSG_BODY(&m, offset, void), Ev, Ev_len);
	offset += Ev_len;
	// copy Dv(Dc(h(Ev))) into M_3
	*MSG_BODY(&m, offset, int) = BN_bn2bin(R3, MSG_BODY(&m, offset+4, unsigned char));

	// send M_3, CTF's response in r
	if (send_recv(&m, &r) == 0){
		assert(MSG_TYPE(&r) == MSG_VOTING_KEY_REPL);
		// did the CTF accept our key?
		// for now when there are conflicts CTF will block all conflicting keys.
		// CTF cannot just allow us to try again since this allows voters to register more
		// than one tallying key and thus vote more than once.
		// there are many ways in which we can ensure that tallying keys are unique
		// for instance we can have each voter pick a unique prime for his public exponent
		// CTF can allow a voter to register in case of a conflict by first asking him to reveal Dv
		if (*MSG_BODY(&r, 4, int) == 1)
			printf("Voting Key Accepted.\n");
		else 
			printf("Voting Key Rejected.\n");
		MSG_FREE(&r);
	}
	else {
		printf("ERROR: 2(e).\n");
		abort();
	}
	MSG_FREE(&m);

	BN_free(R2);
	BN_free(R3);
	//MyPub_RSA = RSAPublicKey_dup(My_RSA);
	//MyPriv_RSA = RSAPrivateKey_dup(My_RSA);
}


// choose IV, K and send encrypted ballot to CTF
void
vote(void)
{
	unsigned char IVKvB[32];

//	assert(RAND_bytes((unsigned char *)&B[0], 8) == 1);
//	B[0] %= 1000;
	B[1] = 0;
	printf("B= %ld\n", B[0]);

	// choose secret key K, and IV, we choose a new secret key, IV
	// each time we vote
	assert(RAND_bytes(K, 16) == 1);
	assert(RAND_bytes(IVKvB, 16) == 1);
	K_valid = true;

	// use 128 bit AES in cbc mode to encrypt ballot
	AES_KEY Ke;
	assert(AES_set_encrypt_key(K, 128, &Ke) == 0);
	unsigned char IV[16];
	memmove(IV, IVKvB, 16);
	// ciphertext is in IV,Kv(B)
	AES_cbc_encrypt((unsigned char *)B, (unsigned char *)&IVKvB[16], 16, &Ke, IV, AES_ENCRYPT);
	
	// compute h(IVKvB) so that we can sign it
	unsigned char hIVKvB[SHA_DIGEST_LENGTH];
	assert(SHA1((unsigned char *)IVKvB, 32, hIVKvB));
	// X is hIVKvB as BIGNUM
	BIGNUM *X = BN_bin2bn(hIVKvB, sizeof(hIVKvB), NULL);
	BIGNUM *R = BN_new();
	// R is Dv(h(IV,Kv(B)))
	assert(BN_mod_exp(R, X, MyPriv_RSA->d, MyPriv_RSA->n, bn_ctx) == 1);

	// create M_4, M_4 has (Ev, IV,Kv(B), Dv(h(IV,Kv(B))))
	MSG m;
	int len2 = BN_num_bytes(R);
	alloc_msg(4+4+Ev_len+32+4+len2, &m);
	MSG_TYPE(&m) = MSG_VOTE;
	int offset = 4;
	*MSG_BODY(&m, offset, int) = Ev_len;
	offset += sizeof(int);
	// copy Ev to M_4
	memmove(MSG_BODY(&m, offset, void), Ev, Ev_len);
	offset += Ev_len;
	// copy IV,Kv(B) to M_4
	memmove(MSG_BODY(&m, offset, void), IVKvB, 32);
	offset += 32;
	*MSG_BODY(&m, offset, int) = len2;
	offset += sizeof(int);
	// copy Dv(h(IV,Kv(B))) to M_4
	assert(len2 == BN_bn2bin(R, MSG_BODY(&m, offset, unsigned char)));
	// send M_4 to CTF. the only reason why CTF may want to reject the ballot is
	// if the checksum (i.e. hash) would not match. we do not handle that for now.
	if (send_recv(&m, NULL) == 0)
		printf("Vote/IV sent.\n");
	else
		printf("ERROR sending vote.\n");
	MSG_FREE(&m);

	BN_free(X);
	BN_free(R);
}


// send secret key K to CTF so that it can decrypt the ballot
void
decrypt_vote()
{
	// did we vote?
	// since client is restarted for each one of these functions
	// we need to check that user is has not made the mistake
	// of asking us to send a secret key before we have computed it.
	// we compute the secret key when we vote.
	if (!K_valid){
		printf("ERROR!!! K is invalid.\n");
		abort();
	}

	// hash the secret key K so that we can sign it with voter's
	// private key
	unsigned char hK[SHA_DIGEST_LENGTH];
	assert(SHA1((unsigned char *)K, sizeof(K), hK));
	// X is h(K)
	BIGNUM *X = BN_bin2bn(hK, sizeof(hK), NULL);
	BIGNUM *R = BN_new();
	// R is Dv(h(K))
	assert(BN_mod_exp(R, X, MyPriv_RSA->d, MyPriv_RSA->n, bn_ctx) == 1);

	// construct M_5 (see doc.)
	// M_5 has format (Ev, K, Dv(h(K)))
	MSG m;
	int len2 = BN_num_bytes(R);
	alloc_msg(4+4+Ev_len+sizeof(K)+4+len2, &m);
	MSG_TYPE(&m) = MSG_DECRYPT_VOTE;
	int offset = 4;
	*MSG_BODY(&m, offset, int) = Ev_len;
	offset += sizeof(int);
	// copy tallying key Ev
	memmove(MSG_BODY(&m, offset, void), Ev, Ev_len);
	offset += Ev_len;
	// copy secret key K
	memmove(MSG_BODY(&m, offset, void), K, sizeof(K));
	offset += sizeof(K);
	*MSG_BODY(&m, offset, int) = len2;
	offset += sizeof(int);
	// copy voter's signature Dv(h(K))
	assert(len2 == BN_bn2bin(R, MSG_BODY(&m, offset, unsigned char)));
	// send M_5 to CTF, as with M_4 we do not currently have
	// CTF acknowledge M_4
	if (send_recv(&m, NULL) == 0)
		printf("K sent.\n");
	else
		printf("ERROR sending K.\n");
	MSG_FREE(&m);

	BN_free(X);
	BN_free(R);
}

int main(int argc, char *argv[])
{
	ERR_load_crypto_strings();
	bn_ctx = BN_CTX_new();

	while (1){
		static struct option long_options[] =
		{
			{"reg", no_argument, &action, REG},
			{"tally-key", no_argument, &action, TALLY_KEY},
			{"tally", no_argument, &action, TALLY},
			{"vote", required_argument, 0, 'v'},
			{"keyfile", required_argument, 0, 'k'},
			{"uid", required_argument, 0, 'u'},
			{"passwd", required_argument, 0, 'p'},
			{0, 0, 0, 0}
		};
		int option_index = 0;
		int c = getopt_long(argc, argv, "",
			long_options, &option_index);
		/*Detect the end of the options.*/
		if(c == -1)
			break;
		switch (c){
			case 'k':
				keyfile = strdup(optarg);
				break;

			case 'v':
				action = VOTE;
				B[0] = atoi(optarg);
				break;

			case 'u':
				uid = atoi(optarg);
				break;

			case 'p':
				passwd = strdup(optarg);
				break;

			case '?':
			case 0:
				break;

			default:
				abort();
		}
	}

	get_CTF_Key();
	switch (action){
		case REG:
			register_to_vote();
			write_keys();
			break;
		case TALLY_KEY:
			read_keys();
			submit_tally_key();
			break;
		case VOTE:
			read_keys();
			vote();
			write_keys();
			break;
		case TALLY:
			read_keys();
			decrypt_vote();
			break;
		default:
			abort();
	}
	return 0;
}
